import cv2

# include image suffixes
IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm'
# include video suffixes
VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv'


class LoadSource:
    def __init__(self, source, img_size=640, stride=32, auto=True) -> None:
        self.source = str(source)
        self.img_size = img_size
        self.stride = stride
        self.auto = auto
        self.running = True

    def stop(self):
        self.running = False

    def __iter__(self):
        suffix = self.source.split('.')[-1].lower()
        # 迭代图片
        if suffix in IMG_FORMATS:
            im0 = cv2.imread(self.source)
            yield im0
        # 视频或者摄像头
        elif suffix in VID_FORMATS or self.source == '0':
            fps = 30
            if self.source == '0':  # 摄像头
                cap = cv2.VideoCapture(int(self.source), cv2.CAP_DSHOW)                
            else:
                cap = cv2.VideoCapture(self.source)
                # fps = cap.get(cv2.CAP_PROP_FPS)
            cap.set(cv2.CAP_PROP_FPS, 30)
            while self.running:
                ret, im0 = cap.read()
                if ret:
                    # time.sleep(1 / fps)
                    im0 = cv2.flip(im0,1)
                    yield im0
                else:
                    print('no frame')
                    break
            cap.release()

        self.running = False
